package com.blue.mountain.mcpsdkjavaserversample;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.node.ArrayNode;
import com.fasterxml.jackson.databind.node.ObjectNode;
import io.modelcontextprotocol.client.McpSyncClient;
import io.modelcontextprotocol.server.McpServer;
import io.modelcontextprotocol.server.McpServerFeatures;
import io.modelcontextprotocol.server.McpSyncServer;
import io.modelcontextprotocol.server.transport.StdioServerTransportProvider;
import io.modelcontextprotocol.spec.McpSchema;
import reactor.core.publisher.Mono;

import java.io.UnsupportedEncodingException;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Map;

public class McpAsyncServerSample {


    public static void main(String[] args) {
        // 设置系统默认编码为UTF-8，确保正确处理中文字符
        System.setProperty("file.encoding", "UTF-8");

        // 创建基于标准输入输出的传输提供者，用于服务器通信
        StdioServerTransportProvider transportProvider = new StdioServerTransportProvider(new ObjectMapper());

        // 创建并配置MCP同步服务器
        McpSyncServer syncServer = McpServer.sync(transportProvider)
                .serverInfo("my-server", "1.0.0") // 设置服务器标识和版本
                .capabilities(McpSchema.ServerCapabilities.builder()
                        .resources(true, true) // 启用资源读写功能
                        .tools(true)          // 启用工具功能
                        .prompts(true)        // 启用提示功能
                        .logging()            // 启用日志功能
                        .build())
                .build();

        try {
            // 添加工具、资源和提示
            addCalculatorTool(syncServer);
            addExampleResource(syncServer);
            addGreetingPrompt(syncServer);

            // 发送日志通知
            syncServer.loggingNotification(McpSchema.LoggingMessageNotification.builder()
                    .level(McpSchema.LoggingLevel.INFO)
                    .logger("custom-logger")
                    .data("Server initialized")
                    .build());

            // 关闭服务器
            // syncServer.close();
        } catch (JsonProcessingException e) {
            // 处理异常，例如打印错误信息
            System.err.println("创建JSON Schema时发生错误: " + e.getMessage());
        }
    }

    private static void addGreetingPrompt(McpSyncServer syncServer) {
        // 创建提示规范，包含提示定义和处理逻辑
        McpServerFeatures.SyncPromptSpecification syncPromptSpecification = new McpServerFeatures.SyncPromptSpecification(
                // 定义提示的基本信息和参数
                new McpSchema.Prompt("greeting", "生成问候语", new ArrayList<>(List.of(
                        new McpSchema.PromptArgument("name", "用户名称", true)))),
                (exchange, request) -> {
                    // 处理提示请求
                    List<McpSchema.PromptMessage> messages = new ArrayList<>();

                    // 处理用户名称，确保正确处理UTF-8编码
                    String name;
                    name = new String(((String) request.arguments().get("name")).getBytes(StandardCharsets.UTF_8), StandardCharsets.UTF_8);
                    if (name.isEmpty()) {
                        name = "访客";
                    }
                    // 创建对话消息序列
                    McpSchema.PromptMessage userMessage = new McpSchema.PromptMessage(
                            McpSchema.Role.USER,
                            new McpSchema.TextContent("你好，请给我一个友好的问候"));

                    McpSchema.PromptMessage systemMessage = new McpSchema.PromptMessage(
                            McpSchema.Role.ASSISTANT,
                            new McpSchema.TextContent("你是一个友好的助手，请为用户生成问候语"));

                    McpSchema.PromptMessage assistantMessage = new McpSchema.PromptMessage(
                            McpSchema.Role.ASSISTANT,
                            new McpSchema.TextContent("你好，" + name + "！很高兴见到你。今天过得怎么样？"));

                    messages.add(systemMessage);
                    messages.add(userMessage);
                    messages.add(assistantMessage);

                    return new McpSchema.GetPromptResult("为用户" + name + "生成的问候语", messages);
                });
        syncServer.addPrompt(syncPromptSpecification);
    }


    /**
     * 创建计算器工具的JSON Schema
     * @return JSON Schema字符串
     * @throws JsonProcessingException JSON处理异常
     */
    private static String createSchema() throws JsonProcessingException {
        ObjectMapper mapper = new ObjectMapper();

        // 创建Schema根节点
        ObjectNode rootNode = mapper.createObjectNode();
        rootNode.put("type", "object");
        rootNode.put("id", "urn:jsonschema:Operation");

        // 创建属性节点
        ObjectNode propertiesNode = mapper.createObjectNode();

        // 定义operation属性：指定可用的计算操作
        ObjectNode operationNode = mapper.createObjectNode();
        operationNode.put("type", "string");
        operationNode.put("description", "计算操作类型，只能是add、subtract、multiply、divide四种");
        ArrayNode enumNode = mapper.createArrayNode();
        enumNode.add("add").add("subtract").add("multiply").add("divide");
        operationNode.set("enum", enumNode);

        // a属性
        ObjectNode aNode = mapper.createObjectNode();
        aNode.put("type", "number");

        // b属性
        ObjectNode bNode = mapper.createObjectNode();
        bNode.put("type", "number");

        // 添加所有属性到properties节点
        propertiesNode.set("operation", operationNode);
        propertiesNode.set("a", aNode);
        propertiesNode.set("b", bNode);

        // 添加properties到根节点
        rootNode.set("properties", propertiesNode);

        // 将JSON对象转换为字符串
        return mapper.writerWithDefaultPrettyPrinter().writeValueAsString(rootNode);
    }

    /**
     * 添加计算器工具到服务器
     * @param syncServer MCP同步服务器实例
     * @throws JsonProcessingException JSON处理异常
     */
    private static void addCalculatorTool(McpSyncServer syncServer) throws JsonProcessingException {
        // 创建工具的JSON Schema定义
        String schema = createSchema();

        // 创建工具规范，包含工具定义和处理逻辑
        McpServerFeatures.SyncToolSpecification syncToolSpecification = new McpServerFeatures.SyncToolSpecification(
                new McpSchema.Tool("calculator", "Basic calculator", schema),
                (exchange, arguments) -> {
                    List<McpSchema.Content> result = new ArrayList<>();

                    try {
                        // 从参数中提取操作类型和操作数
                        String operation = (String) arguments.get("operation");
                        double a = ((Number) arguments.get("a")).doubleValue();
                        double b = ((Number) arguments.get("b")).doubleValue();

                        // 根据操作类型执行相应的计算
                        double calculationResult = 0;
                        switch (operation) {
                            case "add":
                                calculationResult = a + b;
                                break;
                            case "subtract":
                                calculationResult = a - b;
                                break;
                            case "multiply":
                                calculationResult = a * b;
                                break;
                            case "divide":
                                if (b == 0) {
                                    result.add(new McpSchema.TextContent("错误：除数不能为零"));
                                    return new McpSchema.CallToolResult(result, true);
                                }
                                calculationResult = a / b;
                                break;
                            default:
                                result.add(new McpSchema.TextContent("不支持的操作: " + operation));
                                return new McpSchema.CallToolResult(result, true);
                        }

                        result.add(new McpSchema.TextContent("计算结果: " + calculationResult));
                    } catch (Exception e) {
                        // 处理计算过程中的异常
                        result.add(new McpSchema.TextContent("计算错误: " + e.getMessage()));
                        return new McpSchema.CallToolResult(result, true);
                    }

                    return new McpSchema.CallToolResult(result, false);
                });
        syncServer.addTool(syncToolSpecification);
    }

    /**
     * 添加示例资源到服务器
     * @param syncServer MCP同步服务器实例
     */
    private static void addExampleResource(McpSyncServer syncServer) {
        // 创建资源规范，包含资源定义和访问处理逻辑
        McpServerFeatures.SyncResourceSpecification syncResourceSpecification = new McpServerFeatures.SyncResourceSpecification(
                // 定义资源的基本信息：URI、名称、描述和内容类型
                new McpSchema.Resource("custom://resource", "示例资源", "这是一个示例资源", "text/plain", null),
                (exchange, request) -> {
                    // 处理资源读取请求
                    List<McpSchema.ResourceContents> contents = new ArrayList<>();
                    String content = "这是资源内容示例";
                    contents.add(new McpSchema.TextResourceContents("custom://resource/content", "text/plain", content));
                    return new McpSchema.ReadResourceResult(contents);
                });
        syncServer.addResource(syncResourceSpecification);
    }

}
